#include <windows.h>
#include <wtsapi32.h>
#include "PowerControl.h"

__declspec(dllexport) BOOL SystemShutdown(MyShutdownType type, DWORD timeout, LPCWSTR message)
{
    // 获取关机特权
    HANDLE hToken;
    TOKEN_PRIVILEGES tkp;
    
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
    {
        return FALSE;
    }
    
    LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid);
    tkp.PrivilegeCount = 1;
    tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    
    AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, NULL, 0);
    CloseHandle(hToken);
    
    // 设置关机标志
    DWORD flags = 0;
    switch (type)
    {
    case MY_SHUTDOWN_NORMAL:
        flags = EWX_SHUTDOWN;
        break;
    case MY_SHUTDOWN_FORCE:
        flags = EWX_SHUTDOWN | EWX_FORCE;
        break;
    case MY_SHUTDOWN_REBOOT:
        flags = EWX_REBOOT;
        break;
    case MY_SHUTDOWN_LOGOFF:
        flags = EWX_LOGOFF;
        break;
    case MY_SHUTDOWN_HYBRID:
        flags = EWX_SHUTDOWN | EWX_HYBRID_SHUTDOWN;
        break;
    default:
        return FALSE;
    }
    
    // 如果需要，通知所有会话
    if (type != MY_SHUTDOWN_LOGOFF)
    {
        // 创建可修改的字符串缓冲区
        WCHAR title[32] = L"系统关机";
        WCHAR defaultMsg[64] = L"系统即将关机";
        WCHAR* dynamicMsg = NULL;
        
        // 准备消息内容
        if (message)
        {
            size_t len = wcslen(message) + 1;
            dynamicMsg = (WCHAR*)malloc(len * sizeof(WCHAR));
            if (dynamicMsg)
            {
                wcscpy(dynamicMsg, message);
            }
        }
        
        // 发送消息
        WTSSendMessageW(WTS_CURRENT_SERVER_HANDLE, 
                      WTSGetActiveConsoleSessionId(),
                      title, 
                      (DWORD)wcslen(title),
                      dynamicMsg ? dynamicMsg : defaultMsg,
                      dynamicMsg ? (DWORD)wcslen(dynamicMsg) : (DWORD)wcslen(defaultMsg),
                      0, 
                      timeout * 1000, 
                      NULL, 
                      FALSE);
        
        // 释放内存
        if (dynamicMsg)
        {
            free(dynamicMsg);
        }
    }
    
    // 执行关机
    return ExitWindowsEx(flags, SHTDN_REASON_MAJOR_OTHER | SHTDN_REASON_MINOR_OTHER);
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    return TRUE;
}